﻿#include "CHlsToMp4Handler.h"
#include "android_log.h"

#define IO_BUFFER_SIZE 32768

CHlsToMp4Handler::CHlsToMp4Handler()
{
}

CHlsToMp4Handler::~CHlsToMp4Handler()
{
	deinitialize();
}

int CHlsToMp4Handler::initialize(const char* inputFile, const char* outputFile)
{
	int ret = -1;
	do 
	{
		if (inited) {
			break;
		}
		ret = openInput(inputFile);
		if (0 != ret) {
			break;
		}
		writeTemp = new char[tempSize];
		ret = openMp4Output(outputFile);
		if (0 != ret) {
			break;
		}

		workFlag = true;
		inited = true;
	} while (false);
	if (!inited) {
		deinitialize();
	}
	return ret;
}

void CHlsToMp4Handler::deinitialize()
{
	if (workFlag) {
		workFlag = false;
	}
	closeInput();
	closeMp4Output();
	closeOutput();
	if (nullptr != writeTemp) {
		delete[] writeTemp;
		writeTemp = nullptr;
	}
	inited = false;
}

int CHlsToMp4Handler::openInput(const char* url)
{
	int ret = -1;
	do 
	{
		if (nullptr == url) {
			break;
		}
		if ((ret = avformat_open_input(&ifmt_ctx, url, 0, NULL)) < 0) {
			break;
		}
		if ((ret = avformat_find_stream_info(ifmt_ctx, NULL)) < 0) {
			break;
		}
		av_dump_format(ifmt_ctx, 1, url, 0);
		ret = 0;
	} while (false);
	return ret;
}

int CHlsToMp4Handler::openOutput(const char* filePath)
{
	int ret = -1;
	do
	{
		if (nullptr == filePath) {
			break;
		}
		std::string strFile = filePath;
		//fpPosterior = new std::fstream(strFile, std::ios::binary | std::ios::in | std::ios::out);
		fpPosterior = new FileCntl(strFile);
		ret = fpPosterior->openFile(O_RDWR | O_CREAT);
		if (0 != ret) {
			break;
		}
		ret = 0;
	} while (false);
	if (0 != ret) {
		closeInput();
	}
	return ret;
}

int CHlsToMp4Handler::openMp4Output(const char* filePath)
{
	int ret = -1;
	do 
	{
		AVOutputFormat* ofmt = NULL;
		ret = avformat_alloc_output_context2(&ofmt_ctx, nullptr, "mp4", filePath);
		if (ofmt_ctx == nullptr) {
			printf("Can't alloc_output_context ret = %d ,mDestFormat = %s , mDestFilePath = %s \n",
				ret, "mp4", filePath);
			break;
		}
		ofmt_ctx->oformat->flags |= AVFMT_TS_NONSTRICT;
		ofmt_ctx->oformat->flags |= AVFMT_NODIMENSIONS;

		stream_mapping_size = ifmt_ctx->nb_streams;
		stream_mapping = (int*)av_mallocz_array(stream_mapping_size, sizeof(*stream_mapping));
		if (!stream_mapping) {
			av_freep(&stream_mapping);
			break;
		}
		ofmt = ofmt_ctx->oformat;

		for (int i = 0; i < ifmt_ctx->nb_streams; i++) {
			AVStream* out_stream;
			AVStream* in_stream = ifmt_ctx->streams[i];
			AVCodecParameters* in_codecpar = in_stream->codecpar;

			if (in_codecpar->codec_type != AVMEDIA_TYPE_AUDIO &&
				in_codecpar->codec_type != AVMEDIA_TYPE_VIDEO &&
				in_codecpar->codec_type != AVMEDIA_TYPE_SUBTITLE) {
				stream_mapping[i] = -1;
				continue;
			}

			if (in_codecpar->codec_type == AVMEDIA_TYPE_VIDEO) {
				width = in_codecpar->width;
				height = in_codecpar->height;

				//if (width == 0 && height == 0) {
				//	ret = -1;
				//	return ret;
				//}
				videoIndex = in_stream->index;
			}
			if (in_codecpar->codec_type == AVMEDIA_TYPE_AUDIO) {
				audioIndex = in_stream->index;
			}

			stream_mapping[i] = stream_index++;

			out_stream = avformat_new_stream(ofmt_ctx, NULL);
			if (!out_stream) {
				ret = -1;
				return ret;
			}

			ret = avcodec_parameters_copy(out_stream->codecpar, in_codecpar);
			if (ret < 0) {
				return ret;
			}
			out_stream->codecpar->codec_tag = 0;
		}
		if (-1 == videoIndex && -1 == audioIndex) {
			ret = -1;
			break;
		}
		av_dump_format(ofmt_ctx, 0, filePath, 1);

		if (0 != openOutput(filePath)) {
			break;
		}

		avio_buffer = (uint8_t*)av_malloc(IO_BUFFER_SIZE);
		ofmt_ctx->pb = avio_alloc_context(avio_buffer, IO_BUFFER_SIZE,
			AVIO_FLAG_WRITE, this,
			nullptr, io_write, io_seek);
		ofmt_ctx->pb->write_data_type = io_write_data_type;

		ret = avformat_write_header(ofmt_ctx, NULL);
		if (ret < 0) {
			//LOGE("Error occurred when opening output file, ret=%d\n", ret);
			//LOGE("Error occurred: %s\n, width=%d, height=%d", av_err2str(ret), width, height);
			/* close output */
			av_freep(&stream_mapping);
			break;
		}

	} while (false);
	return ret;
}

void CHlsToMp4Handler::closeInput()
{
	if (ifmt_ctx) {
		avformat_close_input(&ifmt_ctx);
		ifmt_ctx = NULL;
	}
}

void CHlsToMp4Handler::closeOutput()
{
	if (fpPosterior) {
		fpPosterior->closeFile();
		delete fpPosterior;
		fpPosterior = nullptr;
	}
}

void CHlsToMp4Handler::closeMp4Output()
{
	if (ofmt_ctx) {
		avformat_free_context(ofmt_ctx);
		ofmt_ctx = NULL;
	}
	if (stream_mapping) {
		av_freep(&stream_mapping);
		stream_mapping = NULL;
	}
}

int CHlsToMp4Handler::run()
{
	if (!workFlag) {
		return -1;
	}
	int ret = -1;
	int64_t video_dts_offset = 0;
	int64_t audio_pts_offset = 0;
	int64_t video_dts_spacing = 0;
	int64_t audio_pts_spacing = 0;
	int64_t last_video_dts = std::numeric_limits<int64_t>::min();
	int64_t last_audio_pts = std::numeric_limits<int64_t>::min();
    int logcount = 0;
	while (workFlag)
	{
		AVStream* in_stream, * out_stream;

		ret = av_read_frame(ifmt_ctx, &pkt);
		if (ret < 0)
			break;
//        if (logcount < 100) {
//            LOGI("read %s pts: %lld, dts:%lld\n", pkt.stream_index == videoIndex ? "video" : "audio", pkt.pts, pkt.dts);
//            logcount ++;
//        }

        in_stream = ifmt_ctx->streams[pkt.stream_index];
		if (pkt.stream_index >= stream_mapping_size ||
			stream_mapping[pkt.stream_index] < 0) {
			av_packet_unref(&pkt);
			continue;
		}

		pkt.stream_index = stream_mapping[pkt.stream_index];
		out_stream = ofmt_ctx->streams[pkt.stream_index];

		if (pkt.pts == AV_NOPTS_VALUE) {
			if (pkt.dts != AV_NOPTS_VALUE) {
				pkt.pts = pkt.dts;
				last_dts = pkt.dts;
			}
			else {
				pkt.pts = last_dts + 1;
				pkt.dts = pkt.pts;
				last_dts = pkt.pts;
			}
		}
		else {
			if (pkt.dts != AV_NOPTS_VALUE) {
				last_dts = pkt.dts;
			}
			else {
				pkt.dts = pkt.pts;
				last_dts = pkt.dts;
			}
		}

		if (pkt.pts < pkt.dts) {
			pkt.pts = pkt.dts;
		}

		//起始的时间归0
//		if (!first_pkt && pkt.flags & AV_PKT_FLAG_KEY) {

		if (!first_pkt) {
			if (!(pkt.flags & AV_PKT_FLAG_KEY)){
				av_packet_unref(&pkt);
				continue;
			}
			first_pts_dts = pkt.pts;
			first_pkt++;
		}

		pkt.pts = pkt.pts - first_pts_dts;
		pkt.dts = pkt.dts - first_pts_dts;
		if (in_stream->codecpar->codec_type == AVMEDIA_TYPE_VIDEO) {
			if (last_video_dts != std::numeric_limits<int64_t>::min() && 0 == video_dts_spacing) {
				video_dts_spacing = pkt.dts - last_video_dts;
			}
			if (0 != video_dts_spacing) {
				if (pkt.dts < last_video_dts) {
					int oldDts = pkt.dts;
					pkt.dts = last_video_dts + video_dts_spacing;
					pkt.pts = pkt.pts + (pkt.dts - oldDts);
				}
				if (pkt.dts - last_video_dts > 2 * video_dts_spacing) {
					video_dts_offset += (pkt.dts - last_video_dts - video_dts_spacing);
				}
			}
			last_video_dts = pkt.dts;
			pkt.pts -= video_dts_offset;
			pkt.dts -= video_dts_offset;
		}
		else {
			if (last_audio_pts != std::numeric_limits<int64_t>::min() && 0 == audio_pts_spacing) {
				audio_pts_spacing = pkt.pts - last_audio_pts;
			}
			if (0 != audio_pts_spacing) {
				if (pkt.pts < last_audio_pts) {
					pkt.pts = last_audio_pts + audio_pts_spacing;
				}
				if (pkt.pts - last_audio_pts > 2 * audio_pts_spacing) {
					audio_pts_offset += (pkt.pts - last_audio_pts - audio_pts_spacing);
				}
			}
			last_audio_pts = pkt.pts;
			pkt.pts -= audio_pts_offset;
//			pkt.dts -= audio_pts_offset;
			pkt.dts = pkt.pts;
		}
//        if (logcount < 100) {
//            LOGI("after %s pts: %lld, dts:%lld\n", pkt.stream_index == videoIndex ? "video" : "audio", pkt.pts, pkt.dts);
//        }

		/* copy packet */
		pkt.pts = av_rescale_q_rnd(pkt.pts, in_stream->time_base, out_stream->time_base,
			static_cast<AVRounding>(AV_ROUND_NEAR_INF |
				AV_ROUND_PASS_MINMAX));
		pkt.dts = av_rescale_q_rnd(pkt.dts, in_stream->time_base, out_stream->time_base,
			static_cast<AVRounding>(AV_ROUND_NEAR_INF |
				AV_ROUND_PASS_MINMAX));
		pkt.duration = av_rescale_q(pkt.duration, in_stream->time_base, out_stream->time_base);
		pkt.pos = -1;

		ret = av_interleaved_write_frame(ofmt_ctx, &pkt);

		//temp_packets_count++;
		//current_progress = temp_packets_count * 1.0f * 100 / total_packets_count;
		////防止JNI回调过分频繁,进行一些代码逻辑上的限制
		//if (abs(current_progress - last_progress) > 0.5f || abs(current_progress - 100) < 0.1f) {
		//	last_progress = current_progress;
		//}

		if (ret < 0) {
			printf("write frame fail: %d\n", ret);
			//break;
		}
		av_packet_unref(&pkt);
	}
	ret = av_write_trailer(ofmt_ctx);
	if (validData > 0) {
		write((uint8_t*)writeTemp, validData);
		lastSeekSize += validData;
		validData = 0;
	}
	resizeFile();
	av_freep(&stream_mapping);
	return ret;
}

int CHlsToMp4Handler::io_write(void* opaque, uint8_t* buf, int size)
{
	auto* ffmpegMux = static_cast<CHlsToMp4Handler*>(opaque);
	return ffmpegMux->muxerWrite(buf, size);
}

int64_t CHlsToMp4Handler::io_seek(void* opaque, int64_t offset, int whence)
{
	auto* ffmpegMux = static_cast<CHlsToMp4Handler*>(opaque);
	return ffmpegMux->muxerSeek(offset, whence);
}

int CHlsToMp4Handler::io_write_data_type(void* opaque, uint8_t* buf, int size, enum AVIODataMarkerType type,
	int64_t time)
{
	auto* ffmpegMux = static_cast<CHlsToMp4Handler*>(opaque);
	CHlsToMp4Handler::DataType dataType = ffmpegMux->mapType(type);
	return ffmpegMux->muxerWriteDataType(buf, size, dataType, time);
}

CHlsToMp4Handler::DataType CHlsToMp4Handler::mapType(AVIODataMarkerType type)
{
	if (type == AVIO_DATA_MARKER_SYNC_POINT) {
		return DATA_SYNC_POINT;
	}

	if (type == AVIO_DATA_MARKER_BOUNDARY_POINT) {
		return DATA_BOUNDARY_POINT;
	}

	if (type == AVIO_DATA_MARKER_FLUSH_POINT) {
		return DATA_FLUSH_POINT;
	}

	if (type == AVIO_DATA_MARKER_HEADER) {
		return DATA_HEADER;
	}

	if (type == AVIO_DATA_MARKER_TRAILER) {
		return DATA_TRAILER;
	}

	if (type == AVIO_DATA_MARKER_UNKNOWN) {
		return DATA_UNKNOWN;
	}

	return DATA_UNKNOWN;
}

int64_t CHlsToMp4Handler::muxerSeek(int64_t offset, int whence)
{
	if (offset > lastSeekSize) {
		lastSeekSize = offset;
	}
	if (validData > 0) {
		write((uint8_t*)writeTemp, validData);
		validData = 0;
	}
	if (!fpPosterior->seekFile(offset, whence)) {
		printf("muxerSeek seek fail\n");
		return 0;
	}
	return offset;
}

int CHlsToMp4Handler::muxerWrite(uint8_t* buf, int size)
{
	if (validData + size >= tempSize) {
		write((uint8_t*)writeTemp, validData);
		validData = 0;
	}
	memcpy(writeTemp + validData, buf, size);
	validData += size;
	return size;
}

int CHlsToMp4Handler::write(uint8_t* buf, int size)
{
	if (size != fpPosterior->writeFile(buf, size)) {
		printf("muxerWrite fail\n");
		return 0;
	}
	return size;
}

int CHlsToMp4Handler::muxerWriteDataType(uint8_t* buf, int size, DataType type, int64_t time)
{
	return muxerWrite(buf, size);
	//printf("muxerWriteDataType size: %d, type: %d, time: %lld\n", size, type, time);
	//return size;
}

void CHlsToMp4Handler::resizeFile()
{
	fpPosterior->truncateFile(lastSeekSize);
}
